from typing import Optional, List, Union
import time

class MockUART:
    """
    Mock UART class
    Used for testing UART communication without hardware
    """
    def __init__(self, uart_id: int, baudrate: int = 9600, bits: int = 8, 
                 parity: Optional[int] = None, stop: int = 1, timeout: int = 1000) -> None:
        self.uart_id = uart_id
        self.baudrate = baudrate
        self.bits = bits
        self.parity = parity
        self.stop = stop
        self.timeout = timeout
        self.sent_data: List[bytes] = []
        self.received_buffer: List[bytes] = []
        self.is_connected = True
        print(f"MockUART {uart_id} initialized: {baudrate}bps, {bits}{parity or 'N'}{stop}")

    def write(self, data: Union[str, bytes]) -> int:
        if not self.is_connected:
            raise OSError("UART not connected")
        if isinstance(data, str):
            data = data.encode('utf-8')
        self.sent_data.append(data)
        print(f"UART {self.uart_id} sent: {data} ({len(data)} bytes)")
        if data == b'AT':
            self.queue_receive_data(b'OK')
        elif data.startswith(b'PING'):
            self.queue_receive_data(b'PONG')
        return len(data)

    def read(self, nbytes: Optional[int] = None) -> Optional[bytes]:
        if not self.received_buffer:
            return None
        if nbytes is None:
            data = b''.join(self.received_buffer)
            self.received_buffer.clear()
        else:
            data = b''
            remaining = nbytes
            while self.received_buffer and remaining > 0:
                chunk = self.received_buffer.pop(0)
                if len(chunk) <= remaining:
                    data += chunk
                    remaining -= len(chunk)
                else:
                    data += chunk[:remaining]
                    self.received_buffer.insert(0, chunk[remaining:])
                    remaining = 0
        if data:
            print(f"UART {self.uart_id} received: {data}")
        return data if data else None

    def any(self) -> int:
        return sum(len(chunk) for chunk in self.received_buffer)

    def queue_receive_data(self, data: Union[str, bytes]) -> None:
        if isinstance(data, str):
            data = data.encode('utf-8')
        self.received_buffer.append(data)

    def disconnect(self) -> None:
        self.is_connected = False
        print(f"UART {self.uart_id} disconnected")

    def __repr__(self) -> str:
        return f"MockUART(id={self.uart_id}, baudrate={self.baudrate}, connected={self.is_connected})"

def uart_send_mock(uart: MockUART, data: Union[str, bytes]) -> Optional[int]:
    try:
        return uart.write(data)
    except Exception as e:
        print(f"Mock UART send error: {e}")
        return None

def uart_receive_mock(uart: MockUART, timeout: int = 1000) -> Optional[bytes]:
    try:
        if uart.any():
            return uart.read()
        return None
    except Exception as e:
        print(f"Mock UART receive error: {e}")
        return None

def uart_send_receive_mock(uart: MockUART, data: Union[str, bytes], timeout: int = 1000) -> Optional[bytes]:
    if uart_send_mock(uart, data):
        time.sleep(0.01)
        return uart_receive_mock(uart, timeout)
    return None
